import os
import torch
import pandas as pd
from torch.utils.data import DataLoader
from tqdm import tqdm

from src.config.config import Config
from src.data.dataset import DigitsDataset
from src.models.resnet import DigitsResnet18, DigitsResnet50
from src.models.mobilenet import DigitsMobilenet
from src.utils.metrics import parse_predictions

class Predictor:
    def __init__(self, config):
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # 初始化测试数据加载器
        self.test_loader = DataLoader(
            DigitsDataset(config.data_dir, mode='test', aug=False),
            batch_size=config.batch_size,
            shuffle=False,
            num_workers=8,
            pin_memory=True,
            drop_last=False
        )

    def predict_single(self, model_path, model_name='resnet18'):
        """单模型预测"""
        # 初始化模型
        if model_name == 'resnet18':
            model = DigitsResnet18(self.config.class_num)
        elif model_name == 'resnet50':
            model = DigitsResnet50(self.config.class_num)
        elif model_name == 'mobilenet':
            model = DigitsMobilenet(self.config.class_num)
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        model = model.to(self.device)

        # 加载模型权重
        checkpoint = torch.load(model_path)
        model.load_state_dict(checkpoint['model_state_dict'])
        model.eval()

        results = []
        with torch.no_grad():
            for img, img_names in tqdm(self.test_loader, desc=f'Predicting with {model_name}'):
                img = img.to(self.device)
                pred = model(img)
                pred_labels = parse_predictions(pred)

                results.extend([(name, code) for name, code in zip(img_names, pred_labels)])

        return sorted(results, key=lambda x: x[0])

    def predict_ensemble(self, model_configs):
        """模型融合预测"""
        models = []
        for config in model_configs:
            if config['name'] == 'resnet18':
                model = DigitsResnet18(self.config.class_num)
            elif config['name'] == 'resnet50':
                model = DigitsResnet50(self.config.class_num)
            elif config['name'] == 'mobilenet':
                model = DigitsMobilenet(self.config.class_num)

            checkpoint = torch.load(config['path'])
            model.load_state_dict(checkpoint['model_state_dict'])
            model = model.to(self.device)
            model.eval()
            models.append((model, config.get('weight', 1.0)))

        results = []
        with torch.no_grad():
            for img, img_names in tqdm(self.test_loader, desc='Ensemble predicting'):
                img = img.to(self.device)

                # 融合预测
                ensemble_pred = [0] * 4  # 4个数字位置
                for model, weight in models:
                    pred = model(img)
                    for i, p in enumerate(pred):
                        ensemble_pred[i] = ensemble_pred[i] + weight * p

                pred_labels = parse_predictions(ensemble_pred)
                results.extend([(name, code) for name, code in zip(img_names, pred_labels)])

        return sorted(results, key=lambda x: x[0])

    def save_results(self, results, save_path):
        """保存预测结果"""
        df = pd.DataFrame(results, columns=['file_name', 'file_code'])
        df['file_name'] = df['file_name'].apply(lambda x: os.path.basename(x))
        df.to_csv(save_path, index=False)
        print(f'Results saved to {save_path}')

def main():
    config = Config()
    predictor = Predictor(config)

    # 单模型预测 - 使用您现有的模型文件
    results = predictor.predict_single(
        model_path='./NDataset/checkpoints/epoch-7-acc-71.18.pth',  # 修改为您的模型文件路径
        model_name='resnet18'  # 确保这是正确的模型类型
    )
    predictor.save_results(results, 'predictions_single.csv')



    # 模型融合预测
    # model_configs = [
    #     {
    #         'name': 'resnet18',
    #         'path': os.path.join(config.checkpoints, 'best_resnet18.pth'),
    #         'weight': 0.4
    #     },
    #     {
    #         'name': 'resnet50',
    #         'path': os.path.join(config.checkpoints, 'best_resnet50.pth'),
    #         'weight': 0.3
    #     },
    #     {
    #         'name': 'mobilenet',
    #         'path': os.path.join(config.checkpoints, 'best_mobilenet.pth'),
    #         'weight': 0.3
    #     }
    # ]
    #
    # results = predictor.predict_ensemble(model_configs)
    # predictor.save_results(results, 'predictions_ensemble.csv')

if __name__ == '__main__':
    main()


